from .hourglass import HourglassNet
from .landrmark import LandMarkNet
from .unet import *
from .anglenet import regressor_cnn
from .pose_hrnet import PoseHighResolutionNet
from .seg_hrnet import HighResolutionNet
from .spine_landmark_resnet import SpineLandmarkResNet

_KEY_POINT_META_ARCHITECTURES={"HourglassNet": HourglassNet,
                               "LandMarkNet":LandMarkNet,
                               'HighResolutionNet': PoseHighResolutionNet,
                               'SpineLandmarkResNet': SpineLandmarkResNet}

_ANGLES_META_ARCHITECTURES={"RegressorNet": regressor_cnn}

_SEMATIC_SEG_META_ARCHITECTURES = {'UNet': UNet,
                                   'FrrNet':FrrNet,
                                   'HighResolutionNet':HighResolutionNet}

def build_keypoint_model(cfg):
    meta_arch = _KEY_POINT_META_ARCHITECTURES[cfg.MODEL.KEYPOINT_REGRESSION.MODEL]
    return meta_arch(cfg.MODEL.KEYPOINT_REGRESSION.CHANNELS,
                          cfg.MODEL.KEYPOINT_REGRESSION.KP_NUM,
                          **cfg.MODEL.KEYPOINT_REGRESSION.MODEL_PARAMETER)


def build_sematicseg_model(cfg):
    meta_arch = _SEMATIC_SEG_META_ARCHITECTURES[cfg.MODEL.SEMANTIC_SEG.MODEL]
    return meta_arch(cfg.MODEL.SEMANTIC_SEG.NUM_CLASSES,
                          **cfg.MODEL.SEMANTIC_SEG.MODEL_PARAMETER)


def build_angle_model(cfg):
    meta_arch = _ANGLES_META_ARCHITECTURES[cfg.MODEL.ANGLES_REGRESSION.MODEL]
    return meta_arch(cfg.MODEL.ANGLES_REGRESSION.CHANNELS,
                          cfg.MODEL.ANGLES_REGRESSION.NUM_ANGLES,
                          **cfg.MODEL.ANGLES_REGRESSION.MODEL_PARAMETER)